// jeMap class
// (c) jimon game studio
// Copyright (C) 2006-2009 by Kat'Oun

#ifndef JEH_JEMAP
#define JEH_JEMAP

#include "jeTypes.h"
#include "jeMemory.h"
#include "jeDebug.h"

namespace je
{
	//! jeMap template for associative arrays using a red-black tree
	template <class KeyType, class ValueType>
	class jeMap
	{
		//! red/black tree for jeMap
		template <class KeyTypeRB, class ValueTypeRB>
		class RBTree
		{
		public:

			RBTree(const KeyTypeRB& k, const ValueTypeRB& v)
				: LeftChild(0), RightChild(0), Parent(0), Key(k),
					Value(v), IsRed(true) {}

			void setLeftChild(RBTree* p)
			{
				LeftChild=p;
				if (p)
					p->setParent(this);
			}

			void setRightChild(RBTree* p)
			{
				RightChild=p;
				if (p)
					p->setParent(this);
			}

			void setParent(RBTree* p)		{ Parent=p; }

			void setValue(const ValueTypeRB& v)	{ Value = v; }

			void setRed()			{ IsRed = true; }
			void setBlack()			{ IsRed = false; }

			RBTree* getLeftChild() const	{ return LeftChild; }
			RBTree* getRightChild() const	{ return RightChild; }
			RBTree* getParent() const	{ return Parent; }

			ValueTypeRB getValue() const
			{
				return Value;
			}

			KeyTypeRB getKey() const
			{
				return Key;
			}

			u1 isRoot() const
			{
				return Parent==0;
			}

			u1 isLeftChild() const
			{
				return (Parent != 0) && (Parent->getLeftChild()==this);
			}

			u1 isRightChild() const
			{
				return (Parent!=0) && (Parent->getRightChild()==this);
			}

			u1 isLeaf() const
			{
				return (LeftChild==0) && (RightChild==0);
			}

			unsigned int getLevel() const
			{
				if (isRoot())
					return 1;
				else
					return getParent()->getLevel() + 1;
			}


			u1 isRed() const
			{
				return IsRed;
			}

			u1 isBlack() const
			{
				return !IsRed;
			}

		private:
			RBTree();

			RBTree*		LeftChild;
			RBTree*		RightChild;

			RBTree*		Parent;

			KeyTypeRB	Key;
			ValueTypeRB	Value;

			u1 IsRed;
		};

		public:

		typedef RBTree<KeyType,ValueType> Node;

		//! Normal Iterator
		class Iterator
		{
		public:

			Iterator() : Root(0), Cur(0) {}

			// Constructor(Node*)
			Iterator(Node* root) : Root(root)
			{
				reset();
			}

			// Copy constructor
			Iterator(const Iterator& src) : Root(src.Root), Cur(src.Cur) {}

			void reset(u1 atLowest=true)
			{
				if (atLowest)
					Cur = getMin(Root);
				else
					Cur = getMax(Root);
			}

			u1 atEnd() const
			{
				return Cur==0;
			}

			Node* getNode()
			{
				return Cur;
			}

			Iterator& operator=(const Iterator& src)
			{
				Root = src.Root;
				Cur = src.Cur;
				return (*this);
			}

			void operator++(int)
			{
				inc();
			}

			void operator--(int)
			{
				dec();
			}


			Node* operator -> ()
			{
				return getNode();
			}

			Node& operator* ()
			{
				JEBREAK(atEnd()) // access violation

				return *Cur;
			}

		private:

			Node* getMin(Node* n)
			{
				while(n && n->getLeftChild())
					n = n->getLeftChild();
				return n;
			}

			Node* getMax(Node* n)
			{
				while(n && n->getRightChild())
					n = n->getRightChild();
				return n;
			}

			void inc()
			{
				// Already at end?
				if (Cur==0)
					return;

				if (Cur->getRightChild())
				{
					// If current node has a right child, the next higher node is the
					// node with lowest key beneath the right child.
					Cur = getMin(Cur->getRightChild());
				}
				else if (Cur->isLeftChild())
				{
					// No right child? Well if current node is a left child then
					// the next higher node is the parent
					Cur = Cur->getParent();
				}
				else
				{
					// Current node neither is left child nor has a right child.
					// Ie it is either right child or root
					// The next higher node is the parent of the first non-right
					// child (ie either a left child or the root) up in the
					// hierarchy. Root's parent is 0.
					while(Cur->isRightChild())
						Cur = Cur->getParent();
					Cur = Cur->getParent();
				}
			}

			void dec()
			{
				// Already at end?
				if (Cur==0)
					return;

				if (Cur->getLeftChild())
				{
					// If current node has a left child, the next lower node is the
					// node with highest key beneath the left child.
					Cur = getMax(Cur->getLeftChild());
				}
				else if (Cur->isRightChild())
				{
					// No left child? Well if current node is a right child then
					// the next lower node is the parent
					Cur = Cur->getParent();
				}
				else
				{
					// Current node neither is right child nor has a left child.
					// Ie it is either left child or root
					// The next higher node is the parent of the first non-left
					// child (ie either a right child or the root) up in the
					// hierarchy. Root's parent is 0.

					while(Cur->isLeftChild())
						Cur = Cur->getParent();
					Cur = Cur->getParent();
				}
			}

			Node* Root;
			Node* Cur;
		}; // Iterator



		//! Parent First Iterator.
		/** Traverses the tree from top to bottom. Typical usage is
		when storing the tree structure, because when reading it
		later (and inserting elements) the tree structure will
		be the same. */
		class ParentFirstIterator
		{
		public:


		ParentFirstIterator() : Root(0), Cur(0)
		{
		}


		explicit ParentFirstIterator(Node* root) : Root(root), Cur(0)
		{
			reset();
		}

		void reset()
		{
			Cur = Root;
		}


		u1 atEnd() const
		{
			return Cur==0;
		}

		Node* getNode()
		{
			return Cur;
		}


		ParentFirstIterator& operator=(const ParentFirstIterator& src)
		{
			Root = src.Root;
			Cur = src.Cur;
			return (*this);
		}


		void operator++(int)
		{
			inc();
		}


		Node* operator -> ()
		{
			return getNode();
		}

		Node& operator* ()
		{
			JEBREAK(atEnd()) // access violation

			return *getNode();
		}

		private:

		void inc()
		{
			// Already at end?
			if (Cur==0)
				return;

			// First we try down to the left
			if (Cur->getLeftChild())
			{
				Cur = Cur->getLeftChild();
			}
			else if (Cur->getRightChild())
			{
				// No left child? The we go down to the right.
				Cur = Cur->getRightChild();
			}
			else
			{
				// No children? Move up in the hierarcy until
				// we either reach 0 (and are finished) or
				// find a right uncle.
				while (Cur!=0)
				{
					// But if parent is left child and has a right "uncle" the parent
					// has already been processed but the uncle hasn't. Move to
					// the uncle.
					if (Cur->isLeftChild() && Cur->getParent()->getRightChild())
					{
						Cur = Cur->getParent()->getRightChild();
						return;
					}
					Cur = Cur->getParent();
				}
			}
		}

		Node* Root;
		Node* Cur;

		}; // ParentFirstIterator


		//! Parent Last Iterator
		/** Traverse the tree from bottom to top.
		Typical usage is when deleting all elements in the tree
		because you must delete the children before you delete
		their parent. */
		class ParentLastIterator
		{
		public:

			ParentLastIterator() : Root(0), Cur(0) {}

			explicit ParentLastIterator(Node* root) : Root(root), Cur(0)
			{
				reset();
			}

			void reset()
			{
				Cur = getMin(Root);
			}

			u1 atEnd() const
			{
				return Cur==0;
			}

			Node* getNode()
			{
				return Cur;
			}

			ParentLastIterator& operator=(const ParentLastIterator& src)
			{
				Root = src.Root;
				Cur = src.Cur;
				return (*this);
			}

			void operator++(int)
			{
				inc();
			}

			Node* operator -> ()
			{
				return getNode();
			}

			Node& operator* ()
			{
				JEBREAK(atEnd()) // access violation

				return *getNode();
			}
		private:

			Node* getMin(Node* n)
			{
				while(n!=0 && (n->getLeftChild()!=0 || n->getRightChild()!=0))
				{
					if (n->getLeftChild())
						n = n->getLeftChild();
					else
						n = n->getRightChild();
				}
				return n;
			}

			void inc()
			{
				// Already at end?
				if (Cur==0)
					return;

				// Note: Starting point is the node as far down to the left as possible.

				// If current node has an uncle to the right, go to the
				// node as far down to the left from the uncle as possible
				// else just go up a level to the parent.
				if (Cur->isLeftChild() && Cur->getParent()->getRightChild())
				{
					Cur = getMin(Cur->getParent()->getRightChild());
				}
				else
					Cur = Cur->getParent();
			}

			Node* Root;
			Node* Cur;
		}; // ParentLastIterator


		// AccessClass is a temporary class used with the [] operator.
		// It makes it possible to have different behavior in situations like:
		// myTree["Foo"] = 32;
		// If "Foo" already exists update its value else insert a new element.
		// int i = myTree["Foo"]
		// If "Foo" exists return its value.
		class AccessClass
		{
			// Let map be the only one who can instantiate this class.
			friend class jeMap<KeyType, ValueType>;

		public:

			// Assignment operator. Handles the myTree["Foo"] = 32; situation
			void operator=(const ValueType& value)
			{
				// Just use the Set method, it handles already exist/not exist situation
				Tree.set(Key,value);
			}

			// ValueType operator
			operator ValueType()
			{
				Node* node = Tree.find(Key);

				// Not found
				JEBREAK(node==0) // access violation

				return node->getValue();
			}

		private:

			AccessClass(jeMap& tree, const KeyType& key) : Tree(tree), Key(key) {}

			AccessClass();

			jeMap& Tree;
			const KeyType& Key;
		}; // AccessClass


		// Constructor.
		jeMap() : Root(0), Size(0) {}

		// Destructor
		~jeMap()
		{
			clear();
		}

		//------------------------------
		// Public Commands
		//------------------------------

		//! Inserts a new node into the tree
		/** \param keyNew: the index for this value
		\param v: the value to insert
		\return True if successful, false if it fails (already exists) */
		u1 insert(const KeyType& keyNew, const ValueType& v)
		{
			// First insert node the "usual" way (no fancy balance logic yet)
			Node* newNode;
			JENEW(newNode,Node(keyNew,v))
			if (!insert(newNode))
			{
				JEDELETE(newNode)
				newNode = NULL;
				
				return false;
			}

			// Then attend a balancing party
			while (!newNode->isRoot() && (newNode->getParent()->isRed()))
			{
				if (newNode->getParent()->isLeftChild())
				{
					// If newNode is a left child, get its right 'uncle'
					Node* newNodesUncle = newNode->getParent()->getParent()->getRightChild();
					if ( newNodesUncle!=0 && newNodesUncle->isRed())
					{
						// case 1 - change the colours
						newNode->getParent()->setBlack();
						newNodesUncle->setBlack();
						newNode->getParent()->getParent()->setRed();
						// Move newNode up the tree
						newNode = newNode->getParent()->getParent();
					}
					else
					{
						// newNodesUncle is a black node
						if ( newNode->isRightChild())
						{
							// and newNode is to the right
							// case 2 - move newNode up and rotate
							newNode = newNode->getParent();
							rotateLeft(newNode);
						}
						// case 3
						newNode->getParent()->setBlack();
						newNode->getParent()->getParent()->setRed();
						rotateRight(newNode->getParent()->getParent());
					}
				}
				else
				{
					// If newNode is a right child, get its left 'uncle'
					Node* newNodesUncle = newNode->getParent()->getParent()->getLeftChild();
					if ( newNodesUncle!=0 && newNodesUncle->isRed())
					{
						// case 1 - change the colours
						newNode->getParent()->setBlack();
						newNodesUncle->setBlack();
						newNode->getParent()->getParent()->setRed();
						// Move newNode up the tree
						newNode = newNode->getParent()->getParent();
					}
					else
					{
						// newNodesUncle is a black node
						if (newNode->isLeftChild())
						{
							// and newNode is to the left
							// case 2 - move newNode up and rotate
							newNode = newNode->getParent();
							rotateRight(newNode);
						}
						// case 3
						newNode->getParent()->setBlack();
						newNode->getParent()->getParent()->setRed();
						rotateLeft(newNode->getParent()->getParent());
					}

				}
			}
			// Color the root black
			Root->setBlack();
			return true;
		}

		//! Replaces the value if the key already exists, otherwise inserts a new element.
		/** \param k The index for this value
		\param v The new value of */
		void set(const KeyType& k, const ValueType& v)
		{
			Node* p = find(k);
			if (p)
				p->setValue(v);
			else
				insert(k,v);
		}

		//! Removes a node from the tree and returns it.
		/** The returned node must be deleted by the user
		\param k the key to remove
		\return A pointer to the node, or 0 if not found */
		Node* delink(const KeyType& k)
		{
			Node* p = find(k);
			if (p == 0)
				return 0;

			// Rotate p down to the left until it has no right child, will get there
			// sooner or later.
			while(p->getRightChild())
			{
				// "Pull up my right child and let it knock me down to the left"
				rotateLeft(p);
			}
			// p now has no right child but might have a left child
			Node* left = p->getLeftChild();

			// Let p's parent point to p's child instead of point to p
			if (p->isLeftChild())
				p->getParent()->setLeftChild(left);

			else if (p->isRightChild())
				p->getParent()->setRightChild(left);

			else
			{
				// p has no parent => p is the root.
				// Let the left child be the new root.
				setRoot(left);
			}

			// p is now gone from the tree in the sense that
			// no one is pointing at it, so return it.

			--Size;
			return p;
		}

		//! Removes a node from the tree and deletes it.
		/** \return True if the node was found and deleted */
		u1 remove(const KeyType& k)
		{
			Node* p = find(k);
			if (p == 0)
			{
				return false;
			}

			// Rotate p down to the left until it has no right child, will get there
			// sooner or later.
			while(p->getRightChild())
			{
				// "Pull up my right child and let it knock me down to the left"
				rotateLeft(p);
			}
			// p now has no right child but might have a left child
			Node* left = p->getLeftChild();

			// Let p's parent point to p's child instead of point to p
			if (p->isLeftChild())
				p->getParent()->setLeftChild(left);

			else if (p->isRightChild())
				p->getParent()->setRightChild(left);

			else
			{
				// p has no parent => p is the root.
				// Let the left child be the new root.
				setRoot(left);
			}

			// p is now gone from the tree in the sense that
			// no one is pointing at it. Let's get rid of it.
			JEDELETE(p)

			--Size;
			return true;
		}

		//! Clear the entire tree
		void clear()
		{
			ParentLastIterator i(getParentLastIterator());

			while(!i.atEnd())
			{
				Node* p = i.getNode();
				i++; // Increment it before it is deleted
					// else iterator will get quite confused.
				JEDELETE(p)
			}
			Root = 0;
			Size= 0;
		}

		//! Is the tree empty?
		//! \return Returns true if empty, false if not
		u1 isEmpty() const
		{
			return Root == 0;
		}

		//! Search for a node with the specified key.
		//! \param keyToFind: The key to find
		//! \return Returns 0 if node couldn't be found.
		Node* find(const KeyType& keyToFind) const
		{
			Node* pNode = Root;

			while(pNode!=0)
			{
				KeyType key(pNode->getKey());

				if (keyToFind == key)
					return pNode;
				else if (keyToFind < key)
					pNode = pNode->getLeftChild();
				else //keyToFind > key
					pNode = pNode->getRightChild();
			}

			return 0;
		}

		//! Gets the root element.
		//! \return Returns a pointer to the root node, or
		//! 0 if the tree is empty.
		Node* getRoot() const
		{
			return Root;
		}

		//! Returns the number of nodes in the tree.
		u32 size() const
		{
			return Size;
		}

		//------------------------------
		// Public Iterators
		//------------------------------

		//! Returns an iterator
		Iterator getIterator()
		{
			Iterator it(getRoot());
			return it;
		}
		//! Returns a ParentFirstIterator.
		//! Traverses the tree from top to bottom. Typical usage is
		//! when storing the tree structure, because when reading it
		//! later (and inserting elements) the tree structure will
		//! be the same.
		ParentFirstIterator getParentFirstIterator()
		{
			ParentFirstIterator it(getRoot());
			return it;
		}
		//! Returns a ParentLastIterator to traverse the tree from
		//! bottom to top.
		//! Typical usage is when deleting all elements in the tree
		//! because you must delete the children before you delete
		//! their parent.
		ParentLastIterator getParentLastIterator()
		{
			ParentLastIterator it(getRoot());
			return it;
		}

		//------------------------------
		// Public Operators
		//------------------------------

		//! operator [] for access to elements
		/** for example myMap["key"] */
		AccessClass operator[](const KeyType& k)
		{
			return AccessClass(*this, k);
		}
		private:

		//------------------------------
		// Disabled methods
		//------------------------------
		// Copy constructor and assignment operator deliberately
		// defined but not implemented. The tree should never be
		// copied, pass along references to it instead.
		explicit jeMap(const jeMap& src);
		jeMap& operator = (const jeMap& src);

		//! Set node as new root.
		/** The node will be set to black, otherwise core dumps may arise
		(patch provided by rogerborg).
		\param newRoot Node which will be the new root
		*/
		void setRoot(Node* newRoot)
		{
			Root = newRoot;
			if (Root != 0)
			{
				Root->setParent(0);
				Root->setBlack();
			}
		}

		//! Insert a node into the tree without using any fancy balancing logic.
		/** \return false if that key already exist in the tree. */
		u1 insert(Node* newNode)
		{
			u1 result=true; // Assume success

			if (Root==0)
			{
				setRoot(newNode);
				Size = 1;
			}
			else
			{
				Node* pNode = Root;
				KeyType keyNew = newNode->getKey();
				while (pNode)
				{
					KeyType key(pNode->getKey());

					if (keyNew == key)
					{
						result = false;
						pNode = 0;
					}
					else if (keyNew < key)
					{
						if (pNode->getLeftChild() == 0)
						{
							pNode->setLeftChild(newNode);
							pNode = 0;
						}
						else
							pNode = pNode->getLeftChild();
					}
					else // keyNew > key
					{
						if (pNode->getRightChild()==0)
						{
							pNode->setRightChild(newNode);
							pNode = 0;
						}
						else
						{
							pNode = pNode->getRightChild();
						}
					}
				}

				if (result)
					++Size;
			}

			return result;
		}

		//! Rotate left.
		//! Pull up node's right child and let it knock node down to the left
		void rotateLeft(Node* p)
		{
			Node* right = p->getRightChild();

			p->setRightChild(right->getLeftChild());

			if (p->isLeftChild())
				p->getParent()->setLeftChild(right);
			else if (p->isRightChild())
				p->getParent()->setRightChild(right);
			else
				setRoot(right);

			right->setLeftChild(p);
		}

		//! Rotate right.
		//! Pull up node's left child and let it knock node down to the right
		void rotateRight(Node* p)
		{
			Node* left = p->getLeftChild();

			p->setLeftChild(left->getRightChild());

			if (p->isLeftChild())
				p->getParent()->setLeftChild(left);
			else if (p->isRightChild())
				p->getParent()->setRightChild(left);
			else
				setRoot(left);

			left->setRightChild(p);
		}

		//------------------------------
		// Private Members
		//------------------------------
		Node* Root; // The top node. 0 if empty.
		u32 Size; // Number of nodes in the tree
	};
}

#endif
